import("io", true)
import("CArray")
import("CMatrix")
import("MPI")
start = MPI.Wtime();

NRA = 500;
NCA = 500;
NCB = 500;

printResult = false;

rank = MPI.Rank();
size = MPI.Size();
workers = size - 1;

b = null;

// Fill arrays
if (rank == 0) {
    a = new CMatrix.Double(NRA, NCA);
    b = new CMatrix.Double(NCA, NCB);

    for (var i=0; i<a.length; i++)
        for (j=0; j<a[i].length; j++)
            a[i][j] = parseFloat(i)+parseFloat(j);

    for (var i=0; i<b.length; i++)
        for (var j=0; j<b[i].length; j++)
            b[i][j] = parseFloat(i)*parseFloat(j);

    sending = MPI.Wtime();
}

// Send B to all nodes
b = MPI.Bcast (0, b);

// Send matrix
if (rank == 0) {
    offset = 0;
    for (dst=1; dst<size; dst++) {
        rows = parseInt(NRA / workers);
        if (dst <= NRA % workers)
            rows += 1;
        MPI.Send(rows, dst);
        MPI.Send(a.slice(offset, offset+rows), dst);
        offset += rows;
    }

    // Wait for results
    c = new CMatrix.Double(NRA, NCB);
    curRow = 0;
    for (src=1; src<size; src++) {
        slice = MPI.Recv(src);
        c.Copy (slice, curRow, 0)
        curRow += slice.length;
    }

    if (printResult) {
        for (var i=0; i<c.length; i++) {
            for (var j=0; j<c[i].length; j++)
                printf (c[i][j] + ",")
            print()
        }
        print()
    }

} else {
    // Receive
    rows = MPI.Recv(0);
    a = MPI.Recv(0);

    // Calculate
    c = new CMatrix.Double(rows, NCB);
    for (k=0; k<NCB; k++) {
        for (i=0; i<rows; i++) {
            c[i][k] = 0.0;
            for (j=0; j<NCA; j++)
                c[i][k] += a[i][j] * b[j][k];
        }
    }

    // Send results
    MPI.Send(c, 0);
}

end = MPI.Wtime();
if (rank == 0) {
    print("Filling: " + (sending-start));
    print("Calculating: " + (end-sending));
    print()
}
